from __future__ import print_function, division
import os
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
import warnings
import csv
warnings.filterwarnings('ignore')
def load_dataset(root_dir, train=True):
    images_path = []
    class_id = []
    if train:
        sub_dir = 'train_dr'
    else:
        sub_dir = 'test_dr'
    with open(os.path.join(root_dir, sub_dir + ".csv"), 'r', encoding='utf-8') as f:
        reader = csv.reader(f)
        for item in reader:
            img_name, id = item
            img_path = os.path.join(root_dir, sub_dir, img_name)
            images_path.append(img_path)
            class_id.append(int(id))
    return images_path, class_id
class Data(Dataset):
    def __init__(self,
                 root_dir,
                 train=True,
                 rotate=45,
                 flip=True,
                 random_crop=True,
                 size=512):

        self.root_dir = root_dir
        self.train = train
        self.rotate = rotate
        self.flip = flip
        self.resize = size
        self.random_crop = random_crop
        self.transform_train = transforms.Compose([
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
            transforms.RandomRotation(degrees=45),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomResizedCrop(size=512),
            transforms.ToTensor(),
            # transforms.RandomErasing(p=0.25),
            transforms.Normalize(mean=[0.321, 0.224, 0.161], std=[0.262, 0.183, 0.132]),  # TortuosityLevel
        ])
        self.transform_test = transforms.Compose([
            transforms.Resize((512, 512)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.321, 0.224, 0.161], std=[0.262, 0.183, 0.132]),
        ])
        self.images_path, self.class_id = load_dataset(self.root_dir, self.train)

    def __len__(self):
        return len(self.images_path)

    def __getitem__(self, idx):
        img_path = self.images_path[idx]
        img_id = self.class_id[idx]
        image = Image.open(img_path).convert("RGB")

        if self.train:
            image = self.transform_train(image)
        else:
            image = self.transform_test(image)
        images = {"img": image,
                  }
        img_class = {"img_id"  : img_id,
                     "img_name": self.images_path[idx]}

        return images, img_class
